#!/usr/local/bin/python

import os
import sys
import re
import subprocess
import socket 
import getpass
import inspect

"""
Generates and sends to lustre a batch of scripts.

USAGE:

    lustre_jobs.py <dir-on-lustre> <bedtools-genome>

EXAMPLE:
    
    lustre_jobs.py /lustre/sblab/berald01/projects/20111206_chipseq_foxm1_debbies/chipseq_pipeline hg18
"""
if len(sys.argv) != 3:
    sys.exit("""%s: Please supply exactly two args:
1. The base-directory on lustre e.g. /lustre/sblab/berald01/projects/20111206_chipseq_foxm1_debbies/chipseq_pipeline
2. The path to the bedtools genome to use or a bedtools tag e.g. hg18""" %(inspect.getfile(inspect.currentframe())))

## Directory on lustre where the jobs are to be executed. It *MUST* exist
LUSTRE_BASEDIR=  sys.argv[1] ## '/lustre/sblab/berald01/projects/20111206_chipseq_foxm1_debbies/chipseq_pipeline'
## Bedtools genome to use
BEDTOOLS_GENOME= sys.argv[2] ## '/home/berald01/applications/bedtools/BEDTools-Version-2.14.3/genomes/human.hg18.genome'
SCRIPT_DIR= 'lustre_scripts'

hostname= socket.gethostname()
username= getpass.getuser()
user_at_host= username + '@' + hostname ## Current Username and host e.g. berald01@uk-cri-lsrv04

try:
    os.makedirs(SCRIPT_DIR)
except OSError:
    pass

listfiles= sorted(os.listdir('bam_clean'))
bamfiles= []
for x in listfiles:
    if re.search('\.bam$', x):
        bamfiles.append(x)

for bam in bamfiles:
    bam= os.path.split(bam)[1]
    bam= os.path.join('bam_clean', bam)
    bam_abspath= os.path.abspath(bam)
    scriptname= os.path.join(SCRIPT_DIR, os.path.split(bam)[1]) + '.sh'
    fout= open(scriptname, 'w')
    sh_script= """
#!/bin/bash

set -e

bedtools=/home/berald01/applications/bedtools/BEDTools-Version-2.14.3/bin

cd %(lustre_basedir)s

mkdir -p bam_clean
scp %(user_at_host)s:%(bam_abspath)s bam_clean  ## Copy bam file from server to lustre

$bedtools/genomeCoverageBed -bg -trackline -ibam %(bam)s -g %(bedtools_genome)s | gzip > %(bam)s.bgr.gz 
$bedtools/genomeCoverageBed -max 200       -ibam %(bam)s -g %(bedtools_genome)s > %(bam)s.genomeCoverageBed.hist 

exit

    """ %{'user_at_host': user_at_host, 'lustre_basedir': LUSTRE_BASEDIR, 'bam': bam, 'bam_abspath': bam_abspath, 'bedtools_genome': BEDTOOLS_GENOME}
    fout.write(sh_script)
    fout.close()
    os.chmod(scriptname, 0744)
    
    cmd= 'ssh berald01@uk-cri-lcst01 mkdir -p %s' %(os.path.join(LUSTRE_BASEDIR, SCRIPT_DIR))
    print(cmd)
    p= subprocess.Popen(cmd, shell= True)
    p.wait()
    
    cmd= 'scp %s berald01@uk-cri-lcst01:%s' %(scriptname, os.path.join(LUSTRE_BASEDIR, SCRIPT_DIR))
    print(cmd)
    p= subprocess.Popen(cmd, shell= True)
    p.wait()
    
    cmd= 'ssh berald01@uk-cri-lcst01 bsub -R "rusage[mem=4096]" -n 1 -o %(script)s.log %(script)s' %{'script': os.path.join(LUSTRE_BASEDIR, scriptname)}
    print(cmd)
    p= subprocess.Popen(cmd, shell= True)
